Improve device-init grouped linear module with single grouped weight support by zhongbozhu · Pull Request #3224 · NVIDIA/TransformerEngine

zhongbozhu · 2026-07-20T21:19:52Z

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

zhongbozhu · 2026-07-23T10:33:15Z

        is_grad_enabled = torch.is_grad_enabled()
        num_gemms = self.num_gemms

-        if FP8GlobalStateManager.fp8_graph_capturing():


Note: this code block was deleted because it was duplicated

zhongbozhu · 2026-07-23T10:55:52Z

/te-ci pytorch

greptile-apps · 2026-07-23T11:07:49Z

Greptile Summary

This PR enhances the device-init grouped linear module to support a single GroupedTensor weight (and bias) parameter shared across all experts, instead of requiring N discrete per-expert parameters. It also adds an explicit grouped_gemm_backend argument to GroupedLinear (replacing the deprecated NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM env-var toggle) and extends tex.group_quantize with an output= parameter that enables in-place MXFP8 re-quantization into an existing workspace for CUDA graph replay.

GroupedLinear gains single_grouped_weight and single_grouped_bias constructor flags; both the module-level and ops-level classes route these through dedicated preparation paths while raising a RuntimeError if the native grouped-tensor backend is unavailable.
The CUDA kernel group_quantize_mxfp8 is updated to compute per-member GEMM-swizzled scale bases correctly for uniform vs. non-uniform tensor groups using a new tensor_rows_for_scales variable.
Extensive new tests cover rejection of host-side m_splits, workspace cache stability, CUDA graph replay of the noop path, and numerical equivalence of single vs. discrete parameter layouts.

Confidence Score: 5/5

Safe to merge; all changed paths are well-tested and the core algorithmic changes are consistent with the existing design.

The refactor is architecturally coherent: single-grouped-parameter routing is enforced by hard errors at both construction and forward time, the backward tensor unpacking correctly accounts for 1 vs N weight/bias arguments, and the CUDA kernel change properly scopes per-member scale bases. Test coverage is extensive. The three inline suggestions are style/annotation concerns that do not affect runtime behavior.

No files require special attention beyond the two annotation-level suggestions in grouped_linear.py.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Large refactor adding single_grouped_weight/bias paths; backward return tuple is correctly structured, workspace caching logic is sound, and the out_features fix (size(0) to size(-2)) is necessary for 3D grouped weight shapes.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Refactored _get_bias_tensors into _get_discrete_bias_tensors and _get_packed_bias_tensor; bias data shape change from 1D to 2D for single_grouped_bias is consistent with the discrete path that already used 2D.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Splits GEMM-swizzled scale base computation into SAME_BOTH_DIMS and VARYING_FIRST_DIM branches; tensor_rows_for_scales correctly bounds colwise scale computation per member.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds output= reuse path to group_quantize: validates num_tensors, logical_shape, and quantizer identity before reusing an existing GroupedTensor, falling through to fresh allocation otherwise.
tests/pytorch/test_grouped_linear.py	Adds comprehensive tests for single grouped weight/bias paths; helper has upper-bound cap at SM 11.0 that will skip tests on future GPUs.
tests/pytorch/test_grouped_tensor.py	New test validates in-place group_quantize reuse (pointer stability), noop_flag behavior, and CUDA graph replay.
tests/pytorch/test_grouped_mlp.py	Adds grouped_tensor_path_supported helper with same SM upper-bound issue; new test validates grouped bias aliases registered parameter storage.

_{Reviews (2): Last reviewed commit: "lint" | Re-trigger Greptile}

timmoon10

The biggest change in this PR is that TE is abandoning any attempt to make single_grouped_weight=True a general feature. Things must be exactly right, or we crash. Given how delicate and experimental this feature has been, I'm not opposed.

The second change is that users must opt-in to access the grouped GEMM kernel. This is also reasonable, since it has alignment requirements for m_splits and it's helpful having a way for users to accept that stricter contract.

We are experiencing many test failures. Given that single_grouped_weight is no longer a general feature, I think it's reasonable we move the corresponding tests to test_grouped_linear.py and test_grouped_mlp.py.

timmoon10 · 2026-07-24T00:54:47Z

+                raise ValueError(
+                    "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 "
+                    "tensor, or select grouped_gemm_backend='legacy'."
+                )


I get that the h2d memcpy is suboptimal, but it's trivially easy to handle. Erroring out seems excessively rigid.

Suggested change

raise ValueError(

"The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 "

"tensor, or select grouped_gemm_backend='legacy'."

)

m_splits = m_splits.to(device=device)

We need to handle the d2h case anyways when the user has specified grouped_gemm_backend="grouped_tensor", but it's not supported and we fallback to split-quantize.

But if the alignment is not provided in the first place, converting it to a device tensor also wouldn't work right, I am okay with another alignment check before adding this H2D copy.

timmoon10 · 2026-07-24T01:13:23Z

                       EXPERIMENTAL and subject to change. Gated by the
                       ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var
                       is not set this argument is forced to ``False`` with a warning.
+    grouped_gemm_backend : {"legacy", "grouped_tensor"}, default = None


These are misleading names.

Legacy implies we are deprecating it, but in fact it is the default backend and it is the fallback when the grouped tensor impl is not supported.

As mentioned in https://github.com/NVIDIA/TransformerEngine/pull/3224/changes#r3618050729, it's confusing that we might not actually perform the backend specified by the user. The user might put grouped_gemm_backend="grouped_tensor" and we'll actually do split-quantize because the cuBLAS version is old or something.

Really, this is not a backend config. It's a hint where the user promises to provide m_splits with a certain alignment. I'd suggest changing this to a bool like enable_grouped_tensor_backend. Alternatively if we want to generalize in case we add more backends in the future, we could have a string like enabled_backends (enabled_backends="grouped_tensor,future_backend").

and we'll actually do split-quantize because the cuBLAS version is old or something

Should we just raise error instead of silently doing this fallback? As a user, it's actually annoying to specify some optimization but it's not working because a fallback get triggered.

I am okay with changing the TE interface into a bool flag like enable_grouped_tensor_backend instead of a str.

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Jul 20, 2026

zhongbozhu force-pushed the improve_device_grouped_linear branch from aa3b9d1 to 47ba66a Compare July 20, 2026 22:36

timmoon10 reviewed Jul 20, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py

zhongbozhu force-pushed the improve_device_grouped_linear branch from ff7eee2 to a43f70f Compare July 20, 2026 22:53

CarlosGomes98 reviewed Jul 21, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py

This was referenced Jul 21, 2026

Single grouped weight fixes #3225

Open

grouped linear single param fixes #3178

Closed

zhongbozhu mentioned this pull request Jul 22, 2026

GroupedLinear(single_grouped_weight=True) produces no weight gradient (wgrad) #3234

Open

zhongbozhu force-pushed the improve_device_grouped_linear branch 2 times, most recently from 6536aa1 to 1169d6e Compare July 23, 2026 10:26

zhongbozhu commented Jul 23, 2026

View reviewed changes

zhongbozhu marked this pull request as ready for review July 23, 2026 10:52

zhongbozhu requested review from Oleg-Goncharov and ksivaman as code owners July 23, 2026 10:52

greptile-apps Bot reviewed Jul 23, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py

Comment thread tests/pytorch/test_grouped_linear.py

Comment thread tests/pytorch/test_grouped_mlp.py

zhongbozhu mentioned this pull request Jul 23, 2026

[Main] Support device-init grouped linear module too with when not using TE opfuser NVIDIA/Megatron-LM#6000

Draft

6 tasks

timmoon10 reviewed Jul 24, 2026

View reviewed changes

zhongbozhu and others added 9 commits July 24, 2026 11:36

improve device grouped linear

68ff7be

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

proper support for grouped bias

e3060fc

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

0450612

for more information, see https://pre-commit.ci

fixes and improvements

2a7a981

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

fix unit test failure

c705d9e

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

fix group linear UT

8187ee0

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

fix group mlp UT by disabling wrong fallback

f1bb519

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

c42ffd9

for more information, see https://pre-commit.ci

lint

69b906f

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

zhongbozhu force-pushed the improve_device_grouped_linear branch from 3d37fd3 to 69b906f Compare July 24, 2026 18:36

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Improve device-init grouped linear module with single grouped weight support #3224

Improve device-init grouped linear module with single grouped weight support #3224
zhongbozhu wants to merge 9 commits into
NVIDIA:mainfrom
zhongbozhu:improve_device_grouped_linear

zhongbozhu commented Jul 20, 2026

Uh oh!

Uh oh!

Uh oh!

zhongbozhu Jul 23, 2026

Uh oh!

zhongbozhu commented Jul 23, 2026

Uh oh!

greptile-apps Bot commented Jul 23, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 left a comment •

edited

Loading

Uh oh!

timmoon10 Jul 24, 2026

Uh oh!

zhongbozhu Jul 24, 2026

Uh oh!

timmoon10 Jul 24, 2026

Uh oh!

zhongbozhu Jul 24, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

zhongbozhu commented Jul 20, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

Uh oh!

Uh oh!

zhongbozhu Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu commented Jul 23, 2026

Uh oh!

greptile-apps Bot commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

timmoon10 Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

timmoon10 Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

greptile-apps Bot commented Jul 23, 2026 •

edited

Loading

timmoon10 left a comment •

edited

Loading

zhongbozhu Jul 24, 2026 •

edited

Loading